// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/core/expressions.h>
#include <pollux/testing/exec/expr_transformer.h>
#include <pollux/type/type.h>

namespace kumo::pollux::exec::test {

using kumo::pollux::TypePtr;
using kumo::pollux::core::TypedExprPtr;
using kumo::pollux::exec::test::ExprTransformer;

class SortArrayTransformer : public ExprTransformer {
 public:
  ~SortArrayTransformer() override = default;

  /// Wraps 'expr' in a call to array_sort. If the type of 'expr' contains a
  /// map, array_sort doesn't support this type, so we return a constant null
  /// instead.
  TypedExprPtr transform(TypedExprPtr expr) const override {
    kumo::pollux::TypePtr type = expr->type();
    if (containsMap(type)) {
      // TODO: support map type by using array_sort with a lambda that casts
      // array elements to JSON before comparison.
      return std::make_shared<kumo::pollux::core::ConstantTypedExpr>(
          type, kumo::pollux::variant::null(type->kind()));
    } else {
      return std::make_shared<kumo::pollux::core::CallTypedExpr>(
          type, std::vector<TypedExprPtr>{std::move(expr)}, "array_sort");
    }
  }

  int32_t extraLevelOfNesting() const override {
    return 1;
  }

 private:
  bool containsMap(const TypePtr& type) const {
    if (type->is_map()) {
      return true;
    } else if (type->is_array()) {
      return containsMap(type->as_array().elementType());
    } else if (type->is_row()) {
      for (const auto& child : type->as_row().children()) {
        if (containsMap(child)) {
          return true;
        }
      }
    }
    return false;
  }
};

} // namespace kumo::pollux::exec::test
